package main

import (
	"regexp"
	"NoFeeProxy/proxy"
	"fmt"
	"flag"
	"net"
	"os"
	"crypto/tls"
	"log"
)

var (
	walletEXP = `"params": \["[\w]{35}"`

	matchid = uint64(0)
	connid  = uint64(0)
	logger  proxy.ColorLogger

	localAddr   = flag.String("l", ":3443", "local address")
	remoteAddr  = flag.String("r", "cn1-zcash.flypool.org:3443", "remote address")
	verbose     = flag.Bool("v", false, "display server actions")
	veryverbose = flag.Bool("vv", false, "display server actions and all tcp data")
	nagles      = flag.Bool("n", false, "disable nagles algorithm")
	hex         = flag.Bool("h", false, "output hex")
	colors      = flag.Bool("c", false, "output ansi colors")
	remoteTLS   = flag.Bool("rtls", true, "remote connection with TLS exposed unencrypted locally")
	localTLS    = flag.Bool("ltls", true, "local connection with TLS exposed unencrypted locally")
	match       = flag.String("match", "", "match regex (in the form 'regex')")
	wallet      = flag.String("w", "t1dFEfbgYinybqjfi3N81ejNKisMDHB6fGC", "your wallet of pool")
	//replace     = flag.String("replace", "", "replace regex (in the form 'regex~replacer')")
)

func main() {
	flag.Parse()

	logger := proxy.ColorLogger{
		Verbose: *verbose,
		Color:   *colors,
	}

	logger.Info("Proxying from %v to %v", *localAddr, *remoteAddr)

	laddr, err := net.ResolveTCPAddr("tcp", *localAddr)
	if err != nil {
		logger.Warn("Failed to resolve local address: %s", err)
		os.Exit(1)
	}
	raddr, err := net.ResolveTCPAddr("tcp", *remoteAddr)
	if err != nil {
		logger.Warn("Failed to resolve remote address: %s", err)
		os.Exit(1)
	}

	var listener *net.TCPListener
	var listenerTls net.Listener
	if *localTLS {
		logger.Info("Starting local tcp connection with TLS")
		crt, err := tls.LoadX509KeyPair("cert/server.crt", "cert/server.key")
		if err != nil {
			log.Fatal(err)
		}
		config := &tls.Config{Certificates: []tls.Certificate{crt}}
		// Time returns the current time as the number of seconds since the epoch.
		// If Time is nil, TLS uses time.Now.
		// Rand provides the source of entropy for nonces and RSA blinding.
		// If Rand is nil, TLS uses the cryptographic random reader in package
		// crypto/rand.
		// The Reader must be safe for use by multiple goroutines.
		listenerTls, err = tls.Listen("tcp", *localAddr, config)
		if err != nil {
			log.Fatal(err)
		}
	} else {
		listener, err = net.ListenTCP("tcp", laddr)
		if err != nil {
			logger.Warn("Failed to open local port to listen: %s", err)
			os.Exit(1)
		}
	}

	matcher := createMatcher(*match)
	replacer := createReplacer(*wallet)

	if *veryverbose {
		*verbose = true
	}

	for {
		var p *proxy.Proxy
		if *remoteTLS {
			logger.Info("Starting remote tcp connection with TLS")
			p = proxy.NewTLSUnwrapped(laddr, raddr, *remoteAddr)
		} else {
			p = proxy.New(laddr, raddr)
		}

		if *localTLS {
			connTls, err := listenerTls.Accept()
			if err != nil {
				log.Panic(err)
			}
			p.SetLconn(connTls)
		} else {
			conn, err := listener.AcceptTCP()
			if err != nil {
				log.Panic(err)
			}
			p.SetLconn(conn)
		}
		if err != nil {
			logger.Warn("Failed to accept connection '%s'", err)
			continue
		}
		connid++

		p.Matcher = matcher
		p.Replacer = replacer

		p.Nagles = *nagles
		p.OutputHex = *hex
		p.Log = proxy.ColorLogger{
			Verbose:     *verbose,
			VeryVerbose: *veryverbose,
			Prefix:      fmt.Sprintf("Connection #%03d ", connid),
			Color:       *colors,
		}

		go p.Start()
	}
}

func createMatcher(match string) func([]byte) {
	if match == "" {
		return nil
	}
	re, err := regexp.Compile(match)
	if err != nil {
		logger.Warn("Invalid match regex: %s", err)
		return nil
	}

	logger.Info("Matching %s", re.String())
	return func(input []byte) {
		ms := re.FindAll(input, -1)
		for _, m := range ms {
			matchid++
			logger.Info("Match #%d: %s", matchid, string(m))
		}
	}
}

func createReplacer(wallet string) func([]byte) []byte {
	if wallet == "" {
		return nil
	}
	re, err := regexp.Compile(walletEXP)
	if err != nil {
		logger.Warn("Invalid replace regex: %s", err)
		return nil
	}

	repl := []byte(fmt.Sprintf(`"params": ["%s.nofee"`, wallet))

	return func(input []byte) []byte {
		if re.Match(input) {
			logger.Info("Replacing %s with %s", re.Find(input), repl)
			return re.ReplaceAll(input, repl)
		}
		return input
	}
}
